from core.agent import Agent

class DontInvestAfterLoss(Agent):
    """Invest if half or more of neighbors invested in the last round."""
    
    def decide(self, my_id, info):
        # In the first round, invest
        if info.current_round == 1:
            return True
        # Subsequently, don 't invest there was a loss in the last round
        return not info.negevent(my_id)

class DontInvestNAfterLoss(Agent):
    """Invest if half or more of neighbors invested in the last round."""
    def decide(self, my_id, info):
        # In the first round, don't invest
        if info.current_round == 1:
            self.dont_invest_count = 0
            return True
        if info.negevent(my_id):
            # invest for 3 rounds after loss
            self.dont_invest_count = 3
            return False
        else:
            self.invest_count =- 1
            return not self.dont_invest_count >= 1
